package com.crscd.ai.config;

import com.crscd.ai.constant.QdrantConstant;
import com.crscd.ai.store.MongoChatMemoryStore;
import com.crscd.ai.store.RedisChatMemoryStore;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections;
import java.util.List;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * Created with IntelliJ IDEA.
 *
 * @author： liuziyang
 * @date： 2025/9/24-19:50
 * @description：
 * @modifiedBy：
 * @version: 1.0
 */
@Slf4j
@Configuration
public class XiaozhiAgentConfig {
  @Autowired private MongoChatMemoryStore mongoChatMemoryStore;
  @Autowired private RedisChatMemoryStore redisChatMemoryStore;

  @Bean
  public ChatMemoryProvider chatMemoryProviderXiaoZhi() {
    return memoryId ->
        MessageWindowChatMemory.builder()
            .id(memoryId)
            .maxMessages(100)
            .chatMemoryStore(redisChatMemoryStore)
            .build();
  }

  @Bean
  public EmbeddingStore<TextSegment> embeddingStore() {
    return QdrantEmbeddingStore.builder()
        .host("10.2.38.103")
        .port(6334)
        .collectionName(QdrantConstant.COLLECTION_NAME)
        .build();
  }

  @Bean
  public ContentRetriever contentRetrieverXiaoZhi(EmbeddingModel embeddingModel) {
    //  从嵌入存储检索和查询内容相关的信息
    return EmbeddingStoreContentRetriever.builder()
        .embeddingStore(embeddingStore())
        .embeddingModel(embeddingModel)
        .maxResults(5)
        .minScore(0.8)
        .build();
  }

  public ContentRetriever contentRetriever(EmbeddingModel embeddingModel) {
    QdrantGrpcClient.Builder grpcClientBuilder =
        QdrantGrpcClient.newBuilder("10.2.38.103", 6334, false);
    final QdrantClient qdrantClient = new QdrantClient(grpcClientBuilder.build());
    // 检查集合是否存在，如果不存在则创建
    try {
      final Boolean exists =
          qdrantClient.collectionExistsAsync(QdrantConstant.COLLECTION_NAME).get();
      if (Boolean.FALSE.equals(exists)) {
        log.info("Qdrant中不存在集合{},下面开始创建这个集合...", QdrantConstant.COLLECTION_NAME);
        // 集合不存在，创建它
        Collections.VectorParams vectorParams =
            Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(QdrantConstant.VECTOR_SIZE)
                .build();
        qdrantClient.createCollectionAsync(QdrantConstant.COLLECTION_NAME, vectorParams);
      }
    } catch (Exception e) {
      log.error("Error checking collection existence", e);
    }

    // 使用FileSystemDocumentLoader读取指定目录下的知识库文档
    // 使用默认的文档解析器对文档进行解析
    Document document1 =
        FileSystemDocumentLoader.loadDocument(
            "/home/lzy/Code/Gitee/javatest/java-ai-langchain4j/src/main/resources/file/hospital_info.md");
    Document document2 =
        FileSystemDocumentLoader.loadDocument(
            "/home/lzy/Code/Gitee/javatest/java-ai-langchain4j/src/main/resources/file/department_info.md");
    Document document3 =
        FileSystemDocumentLoader.loadDocument(
            "/home/lzy/Code/Gitee/javatest/java-ai-langchain4j/src/main/resources/file/neurology_dept_info.md");
    List<Document> documents = List.of(document1, document2, document3);

    // 使用内存向量存储
    //    InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
    // 使用默认的文本分割器
    // 显式创建并指定嵌入模型
    //    EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

    // 使用 Builder 配置 EmbeddingStoreIngestor
    EmbeddingStoreIngestor ingestor =
        EmbeddingStoreIngestor.builder()
            .documentSplitter(DocumentSplitters.recursive(300, 30))
            .embeddingModel(embeddingModel)
            .embeddingStore(embeddingStore())
            .build();

    ingestor.ingest(documents);
    //  从嵌入存储检索和查询内容相关的信息
    return EmbeddingStoreContentRetriever.builder()
        .embeddingStore(embeddingStore())
        .embeddingModel(embeddingModel)
        .maxResults(5)
        .minScore(0.7)
        .build();
  }
}
